package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"mcp-server/internal/config"
	"mcp-server/internal/docprocessor"
	"mcp-server/internal/metadata"
	"mcp-server/internal/qdrant"
	"mcp-server/internal/vector"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Println("使用方法: go run main.go <知识库目录路径>")
		fmt.Println("示例: go run main.go ../../docs/样本知识")
		os.Exit(1)
	}

	knowledgeDir := os.Args[1]
	fmt.Printf("=== 带元数据的知识库向量化处理 ===\n")
	fmt.Printf("处理目录: %s\n\n", knowledgeDir)

	// 加载配置
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatalf("加载配置失败: %v", err)
	}

	// 初始化组件
	docProcessor := docprocessor.NewDocumentProcessor()
	vectorizer := vector.NewOpenAIVectorizer(cfg.OpenAI.APIKey, cfg.OpenAI.BaseURL)
	qdrantClient := qdrant.NewClient(cfg.Qdrant.URL, cfg.Qdrant.APIKey)

	// 创建元数据管理器
	metadataStorageDir := filepath.Join("metadata_storage")
	metadataManager := metadata.NewDefaultManager(metadataStorageDir)

	// 确保集合存在
	collectionName := "knowledge_with_metadata"
	err = qdrantClient.EnsureCollection(context.Background(), collectionName, 1536)
	if err != nil {
		log.Fatalf("创建集合失败: %v", err)
	}

	// 统计信息
	var totalFiles, processedFiles, totalChunks int
	var totalProcessTime time.Duration
	startTime := time.Now()

	// 遍历处理文件
	err = filepath.Walk(knowledgeDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		// 只处理支持的文件类型
		ext := strings.ToLower(filepath.Ext(path))
		if ext != ".md" && ext != ".txt" {
			return nil
		}

		totalFiles++
		fmt.Printf("处理文件 %d: %s\n", totalFiles, path)

		// 1. 提取文档元数据
		fmt.Printf("  1. 提取元数据...")
		docMetadata, err := metadataManager.ProcessFile(path)
		if err != nil {
			fmt.Printf(" 失败: %v\n", err)
			return nil
		}
		fmt.Printf(" 完成 (质量: %.1f, 置信度: %.1f)\n", docMetadata.Quality, docMetadata.Confidence)

		// 2. 文档分片处理
		fmt.Printf("  2. 文档分片...")
		fileStartTime := time.Now()
		result, err := docProcessor.ProcessFile(path)
		if err != nil {
			fmt.Printf(" 失败: %v\n", err)
			return nil
		}
		fileProcessTime := time.Since(fileStartTime)
		totalProcessTime += fileProcessTime
		fmt.Printf(" 完成 (%d个分片, %v)\n", len(result.Chunks), fileProcessTime)

		// 3. 向量化处理
		fmt.Printf("  3. 向量化处理...")
		vectorStartTime := time.Now()

		for i, chunk := range result.Chunks {
			// 创建增强的元数据
			enhancedMetadata := createEnhancedMetadata(chunk, docMetadata, path)

			// 生成向量
			embedding, err := vectorizer.GenerateEmbedding(chunk.Content)
			if err != nil {
				fmt.Printf(" 分片%d向量化失败: %v\n", i+1, err)
				continue
			}

			// 存储到Qdrant
			point := qdrant.Point{
				ID:       chunk.ID,
				Vector:   embedding,
				Metadata: enhancedMetadata,
			}

			err = qdrantClient.UpsertPoints(context.Background(), collectionName, []qdrant.Point{point})
			if err != nil {
				fmt.Printf(" 分片%d存储失败: %v\n", i+1, err)
				continue
			}

			totalChunks++
		}

		vectorTime := time.Since(vectorStartTime)
		fmt.Printf(" 完成 (%d个向量, %v)\n", len(result.Chunks), vectorTime)

		processedFiles++
		fmt.Printf("  ✓ 文件处理完成\n\n")
		return nil
	})

	if err != nil {
		log.Fatalf("遍历目录失败: %v", err)
	}

	// 输出统计信息
	totalTime := time.Since(startTime)
	fmt.Printf("=== 处理完成 ===\n")
	fmt.Printf("总文件数: %d\n", totalFiles)
	fmt.Printf("成功处理: %d\n", processedFiles)
	fmt.Printf("总分片数: %d\n", totalChunks)
	fmt.Printf("总耗时: %v\n", totalTime)
	fmt.Printf("平均每文件: %v\n", totalTime/time.Duration(max(processedFiles, 1)))
	fmt.Printf("文档处理耗时: %v\n", totalProcessTime)
	fmt.Printf("向量化耗时: %v\n", totalTime-totalProcessTime)

	// 展示元数据统计
	fmt.Printf("\n=== 元数据统计 ===\n")
	stats, err := metadataManager.GetStatistics()
	if err != nil {
		fmt.Printf("获取元数据统计失败: %v\n", err)
	} else {
		fmt.Printf("平均质量评分: %.2f\n", stats["average_quality"])
		fmt.Printf("平均置信度: %.2f\n", stats["average_confidence"])
		fmt.Printf("文档分类分布: %v\n", stats["categories"])
		fmt.Printf("语言分布: %v\n", stats["languages"])
	}

	fmt.Printf("\n知识库向量化完成！集合名称: %s\n", collectionName)
}

// createEnhancedMetadata 创建增强的元数据
func createEnhancedMetadata(chunk docprocessor.Chunk, docMeta *metadata.DocumentMetadata, filePath string) map[string]interface{} {
	// 基础分片元数据
	metadata := map[string]interface{}{
		"file_path":      filePath,
		"chunk_id":       chunk.ID,
		"start_line":     chunk.StartLine,
		"end_line":       chunk.EndLine,
		"semantic_tag":   chunk.SemanticTag,
		"content_type":   chunk.ContentType,
		"chunk_index":    chunk.Index,
		"content_length": len([]rune(chunk.Content)),
	}

	// 添加文档级元数据
	if docMeta != nil {
		metadata["doc_title"] = docMeta.Title
		metadata["doc_author"] = docMeta.Author
		metadata["doc_category"] = docMeta.Category
		metadata["doc_language"] = docMeta.Language
		metadata["doc_quality"] = docMeta.Quality
		metadata["doc_confidence"] = docMeta.Confidence
		metadata["doc_word_count"] = docMeta.WordCount
		metadata["doc_created_at"] = docMeta.CreatedAt.Format(time.RFC3339)
		metadata["doc_modified_at"] = docMeta.ModifiedAt.Format(time.RFC3339)

		// 关键词（转换为字符串便于搜索）
		if len(docMeta.Keywords) > 0 {
			metadata["doc_keywords"] = strings.Join(docMeta.Keywords, ",")
		}

		// 标题层级（转换为字符串）
		if len(docMeta.Headings) > 0 {
			metadata["doc_headings"] = strings.Join(docMeta.Headings, "|")
		}

		// 摘要（截取前200字符）
		if docMeta.Summary != "" {
			summary := docMeta.Summary
			if len([]rune(summary)) > 200 {
				summary = string([]rune(summary)[:200]) + "..."
			}
			metadata["doc_summary"] = summary
		}
	}

	// 添加处理时间戳
	metadata["processed_at"] = time.Now().Format(time.RFC3339)
	metadata["metadata_version"] = "1.0"

	return metadata
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}
